'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: file_extractor.py
* @Time: 2025/9/26
* @All Rights Reserve By Brtc
'''
import os
import tempfile
from dataclasses import dataclass
from pathlib import Path
from typing import Union

import requests
from injector import inject
from langchain_community.document_loaders import UnstructuredExcelLoader, UnstructuredPDFLoader, \
    UnstructuredMarkdownLoader, UnstructuredHTMLLoader, UnstructuredCSVLoader, UnstructuredXMLLoader, \
    UnstructuredFileLoader, TextLoader

from internal.model import UploadFile
from internal.service import CosService
from langchain_core.documents import Document as LCDocument


@inject
@dataclass
class FileExtractor:
    """文件提取， 用于将用户上传的文件记录， 加载成Laangchain 对应的Document"""
    cos_service:CosService
    def load(self, upload_file:UploadFile, return_text:bool = False, is_unstructured:bool=False)->Union[list[LCDocument], str]:
        """加载传日的upload_file记录 返回Langchain 文档列表或者字符串 """
        #1、创建一个临时文件夹
        with tempfile.TemporaryDirectory() as temp_dir:
            #2、构建一个临时文件路径
            file_path = os.path.join(temp_dir, os.path.basename(upload_file.key))
            #3、将对象存储中的文件下载到本地
            self.cos_service.download_file(upload_file.key, file_path)
            #4、从指定的路径中去加载文件
            return self.load_from_file(file_path, return_text, is_unstructured)

    @classmethod
    def load_from_file(cls, file_path:str, return_text:bool = False, is_unstructured:bool=False)->Union[list[LCDocument], str]:
        """从本地文件中加载数据， 返回Langchain文档列表"""
        # 1、获取文件扩展名
        delimiter = "\n\n"
        file_extenion = Path(file_path).suffix.lower()

        #2、根据不同的扩展名 去加载不同的加载器
        if file_extenion in [".xlsx",".xls"]:
            loader = UnstructuredExcelLoader(file_path)
        elif file_extenion==".pdf":
            loader = UnstructuredPDFLoader(file_path)
        elif file_extenion in [".md",".markdown"]:
            loader = UnstructuredMarkdownLoader(file_path)
        elif file_extenion in [".htm",".html"]:
            loader = UnstructuredHTMLLoader(file_path)
        elif file_extenion == ".csv":
            loader = UnstructuredCSVLoader(file_path)
        elif file_extenion == ".xml":
            loader = UnstructuredXMLLoader(file_path)
        else:
            loader = UnstructuredFileLoader(file_path) if is_unstructured else TextLoader(file_path)
        return delimiter.join([document.page_content for document in loader.load()]) if return_text else loader.load()

    @classmethod
    def load_from_url(cls, url:str, return_text:bool = False)->Union[list[LCDocument], str]:
        """从传入的远程Url中去加载数据，返回langchain 的文档列表或者字符串"""
        #1、下载远程的文件到本地
        response = requests.get(url)

        with tempfile.TemporaryDirectory() as temp_dir:
            # 2、构建一个临时文件路径
            file_path = os.path.join(temp_dir, os.path.basename(url))
            # 3、将url中的文件下载到本地
            with open(file_path, "wb") as file:
                file.write(response.content)
            # 4、从指定的路径中去加载文件
            return cls.load_from_file(file_path, return_text)






















